# ai_processor.py
import re
import os
from volcenginesdkarkruntime import Ark
from dotenv import load_dotenv

# 加载.env文件（默认会从项目根目录查找）
load_dotenv()

class AIGenerator:
    def __init__(self, engine='volc'):
        self.engine = engine
        self.volc_client = Ark(
            # api_key=os.environ.get("ARK_API_KEY"),
            api_key = os.getenv("ARK_API_KEY"),
            base_url = "https://ark.cn-beijing.volces.com/api/v3"
        )
        # 初始化违禁词列表（建议从文件或数据库读取）
        self.prohibited_words = os.getenv("PROHIBITED_WORDS").split(",")

    def _filter_prohibited_words(self, text):
        """违禁词过滤方法"""
        # 使用正则表达式进行智能匹配
        pattern = re.compile("|".join(
            [re.escape(word) for word in self.prohibited_words]
        ), flags=re.IGNORECASE)
        
        # 替换违禁词为*
        filtered_text = pattern.sub("[内容已过滤]", text)
        
        # 如果发现违禁词可以记录日志或触发警报
        if filtered_text != text:
            print(f"发现违禁词并已过滤，原始内容：{text}")
            
        return filtered_text

    def generate_reply(self, prompt):
        # print(prompt)
        # if self.engine == 'gpt':
        raw_reply = self.volc_client.chat.completions.create(
            model="deepseek-v3-241226",
            messages=[
                {"role": "system", "content": "你是一个活泼的web前端直播助手，请用简短口语化中文回复，不要超过100个字，不要出现敏感词"},
                {"role": "user", "content": prompt},
            ],
            # 响应内容是否流式返回
            # stream=True,
        )

        # 过滤违禁词
        safe_reply = self._filter_prohibited_words(raw_reply.choices[0].message.content)
        return safe_reply
